package com.xnyzc.lhy.common.util;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import com.esotericsoftware.kryo.serializers.CollectionSerializer;
import com.esotericsoftware.kryo.serializers.JavaSerializer;
import com.esotericsoftware.kryo.serializers.MapSerializer;
import com.xnyzc.lhy.common.exception.PangException;
import com.xnyzc.lhy.common.penum.EErrorCode;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.util.*;

/**
 * kryo utils
 */
public class KryoUtil {

    /**
     * 序列化
     *
     * @return
     */
    public static <T> byte[] serializationObjectToByteArr(T object) {
        if (object == null) {
            throw new PangException(EErrorCode.missingArg, "arg error");
        }

        Kryo kryo = new Kryo();
        kryo.setReferences(false);
        kryo.register(object.getClass(), new JavaSerializer());

        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        Output output = new Output(baos);
        kryo.writeObject(output, object);
        output.flush();
        output.close();
        byte[] bys = baos.toByteArray();
        try {
            baos.flush();
            baos.close();
        }
        catch (Exception e) {
            throw new PangException("KryoUtil IO EXCEPTION", e);
        }

        return Base64.getEncoder().encode(bys);
    }

    /**
     * 序列化
     *
     * @return
     */
    public static <T> String serializationObject(T object) {
        return new String(serializationObjectToByteArr(object));
    }

    /**
     * 序列化列表
     *
     * @param obj
     * @param clazz
     * @param <T>
     * @return
     */
    public static <T> String serializationList(List<T> obj, Class<T> clazz) {
        if (obj == null || clazz == null) {
            throw new PangException(EErrorCode.missingArg, "arg error");
        }

        Kryo kryo = new Kryo();
        kryo.setReferences(false);
        kryo.setRegistrationRequired(true);

        CollectionSerializer serializer = new CollectionSerializer();
        serializer.setElementClass(clazz, new JavaSerializer());
        serializer.setElementsCanBeNull(false);

        kryo.register(clazz, new JavaSerializer());
        kryo.register(ArrayList.class, serializer);

        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        Output output = new Output(baos);
        kryo.writeObject(output, obj);
        output.flush();
        output.close();

        byte[] b = baos.toByteArray();
        try {
            baos.flush();
            baos.close();
        }
        catch (IOException e) {
            throw new PangException("KryoUtil IO EXCEPTION", e);
        }

        return new String(Base64.getEncoder().encode(b));
    }

    /**
     * 序列化map
     *
     * @param obj
     * @param clazz
     * @param <T>
     * @return
     */
    public static <T> String serializationMap(Map<String, T> obj, Class<T> clazz) {
        if (obj == null || clazz == null) {
            throw new PangException(EErrorCode.missingArg, "arg error");
        }

        Kryo kryo = new Kryo();
        kryo.setReferences(false);
        kryo.setRegistrationRequired(true);

        MapSerializer serializer = new MapSerializer();
        serializer.setKeyClass(String.class, new JavaSerializer());
        serializer.setKeysCanBeNull(false);
        serializer.setValueClass(clazz, new JavaSerializer());
        serializer.setValuesCanBeNull(true);

        kryo.register(clazz, new JavaSerializer());
        kryo.register(HashMap.class, serializer);

        ByteArrayOutputStream baos = new ByteArrayOutputStream();
        Output output = new Output(baos);
        kryo.writeObject(output, obj);
        output.flush();
        output.close();

        byte[] b = baos.toByteArray();
        try {
            baos.flush();
            baos.close();
        }
        catch (IOException e) {
            throw new PangException("KryoUtil IO EXCEPTION", e);
        }

        return new String(Base64.getEncoder().encode(b));
    }


    /**
     * 反序列化
     *
     * @param clazz
     * @param <T>
     * @return
     */
    public static <T> T deserializationObject(String obj, Class<T> clazz) {
        if (obj == null || clazz == null) {
            throw new PangException(EErrorCode.missingArg, "arg error");
        }

        Kryo kryo = new Kryo();
        kryo.setReferences(false);
        kryo.register(clazz, new JavaSerializer());

        ByteArrayInputStream bais = new ByteArrayInputStream(Base64.getDecoder().decode(obj));
        Input input = new Input(bais);
        return (T) kryo.readClassAndObject(input);
    }

    /**
     * 反序列化列表
     *
     * @param obj
     * @param clazz
     * @param <T>
     * @return
     */
    public static <T> List<T> deserializationList(String obj, Class<T> clazz) {
        if (obj == null || clazz == null) {
            throw new PangException(EErrorCode.missingArg, "arg error");
        }

        Kryo kryo = new Kryo();
        kryo.setReferences(false);
        kryo.setRegistrationRequired(true);

        CollectionSerializer serializer = new CollectionSerializer();
        serializer.setElementClass(clazz, new JavaSerializer());
        serializer.setElementsCanBeNull(false);

        kryo.register(clazz, new JavaSerializer());
        kryo.register(ArrayList.class, serializer);

        ByteArrayInputStream bais = new ByteArrayInputStream(Base64.getDecoder().decode(obj));

        Input input = new Input(bais);
        return (List<T>) kryo.readObject(input, ArrayList.class, serializer);
    }

    /**
     * 反序列化map
     *
     * @param obj
     * @param clazz
     * @param <T>
     * @return
     */
    public static <T> Map<String, T> deserializationMap(String obj, Class<T> clazz) {
        if (obj == null || clazz == null) {
            throw new PangException(EErrorCode.missingArg, "arg error");
        }

        Kryo kryo = new Kryo();
        kryo.setReferences(false);
        kryo.setRegistrationRequired(true);

        MapSerializer serializer = new MapSerializer();
        serializer.setKeyClass(String.class, new JavaSerializer());
        serializer.setKeysCanBeNull(false);
        serializer.setValueClass(clazz, new JavaSerializer());
        serializer.setValuesCanBeNull(true);

        kryo.register(clazz, new JavaSerializer());
        kryo.register(HashMap.class, serializer);

        ByteArrayInputStream bais = new ByteArrayInputStream(Base64.getDecoder().decode(obj));

        Input input = new Input(bais);
        return (Map<String, T>) kryo.readObject(input, HashMap.class, serializer);
    }
}
